'''
Description: hand controller
Version: V1.0
Author: zw_1520@163.com
Date: 2025-06-05 09:38:30
LastEditors: zw_1520@163.com
LastEditTime: 2025-06-05 15:33:41
All rights reserved.
'''


from typing import Text
import sys

import numpy as np
import rclpy
from PySide6.QtCore import Qt, Signal
from PySide6.QtWidgets import (QApplication, QHBoxLayout, QLabel, QSlider,
                               QVBoxLayout, QWidget, QFrame, QLineEdit, QPushButton)
from rclpy.node import Node
from sensor_msgs.msg import JointState

thumb_rotation_value = 0.0
thumb_swing_value = 0.0
index_rotation_value = 0.0
middle_rotation_value = 0.0
ring_rotation_value = 0.0
little_rotation_value = 0.0


class FingerControlWidget(QWidget):
    def __init__(self, node: Node, finger_name: Text, joint_name: Text, hand_name: Text):
        """Finger control widget for Corenetic dexterous hand.

        Args:
            node: ROS2 node
            finger_name: finger name, e.g. thumb, index, middle, ring
            joint_name: joint name, e.g. swing, bend_1, bend_2
            hand_name: hand name, e.g. right, left
        """
        super().__init__()
        self.node = node
        self.finger_name = finger_name
        self.joint_name = joint_name
        self.hand = hand_name

        # Set color based on finger and joint names
        finger_colors = {
            'thumb': 'red',
            'index': 'green',
            'middle': 'blue',
            'ring': 'black'
        }
        joint_colors = {
            'swing': '#FF8C00',  # Dark Orange
            'bend_1': '#00CED1',  # Dark Turquoise
            'bend_2': '#8B008B',  # Dark Magenta
            'bend': '#00CED1'
        }
        finger_color = finger_colors.get(finger_name, 'black')
        joint_color = joint_colors.get(joint_name, 'black')

        self.layout = QHBoxLayout()
        self.label = QLabel(
            f"<span style='color:{finger_color}'>{finger_name}</span> <span style='color:{joint_color}'>{joint_name}</span>")
        # self.label.setStyleSheet(f"color: {finger_color}")
        self.label.setFixedWidth(100)
        self.slider = QSlider(Qt.Horizontal)
        self.slider.setRange(0, 200)
        self.slider.valueChanged.connect(self.on_slider_value_changed)
        self.value_label = QLabel("0")

        self.layout.addWidget(self.label)
        self.layout.addWidget(self.slider)
        self.layout.addWidget(self.value_label)
        self.setLayout(self.layout)

        self.publisher = self.node.create_publisher(
            JointState, f"{self.hand}/target_joint_states", 10)

    def on_slider_value_changed(self, value):
        global thumb_rotation_value
        global index_rotation_value
        global middle_rotation_value
        global ring_rotation_value
        global little_rotation_value
        global thumb_swing_value

        joint_state = JointState()
        # Sets frame_id
        joint_state.header.frame_id = 'rohand_2'

        print(f'Slider value changed to {value}')
        print(f'finger_name: {self.finger_name}')
        # Sets all joints value
        if self.finger_name == 'thumb' and self.joint_name == 'swing':
            joint_state.name = [f"{self.finger_name}"]
            thumb_swing_value = value * 1.0
        elif self.finger_name == 'thumb' and self.joint_name == 'rotation':
            joint_state.name = [f"{self.finger_name}_{self.joint_name}"]
            thumb_rotation_value = value * 1.0
        else:
            joint_state.name = [f"{self.finger_name}"]

        if self.finger_name == 'index':
            joint_state.name = [f"{self.finger_name}_{self.joint_name}"]
            index_rotation_value = value * 1.0
        elif self.finger_name == 'middle':
            joint_state.name = [f"{self.finger_name}_{self.joint_name}"]
            middle_rotation_value = value * 1.0
        elif self.finger_name == 'ring':
            joint_state.name = [f"{self.finger_name}_{self.joint_name}"]
            ring_rotation_value = value * 1.0
        elif self.finger_name == 'little':
            joint_state.name = [f"{self.finger_name}_{self.joint_name}"]
            little_rotation_value = value * 1.0
        # joint_state.name = [f"{self.finger_name}_{self.joint_name}"]
        joint_state.name = ['thumb', 'index', 'middle',
                            'ring', 'little', 'thumb_rotation']
        joint_state.position = [thumb_swing_value, index_rotation_value, middle_rotation_value,
                                ring_rotation_value, little_rotation_value, thumb_rotation_value]
        joint_state.velocity = [65535.0, 65535.0,
                                65535.0, 65535.0, 65535.0, 65535.0]
        joint_state.effort = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        self.value_label.setText(f"{value}")
        print(f'Pub: {joint_state}')
        self.publisher.publish(joint_state)

    def set_value(self, value):
        value = np.round(value * 100).astype(int)
        self.slider.setValue(value)
        self.value_label.setText(str(value))

    def publish_value(self, value):
        joint_state = JointState()
        joint_state.name = [f"{self.finger_name}_{self.joint_name}"]
        joint_state.position = [value]
        joint_state.velocity = [0.0]
        joint_state.effort = [0.0]
        self.value_label.setText(f"{value}")
        print(f'Pub: {joint_state}')
        self.publisher.publish(joint_state)
        value = np.round(value * 100).astype(int)
        self.slider.setValue(value)


class HandControlNode(Node):
    def __init__(self, hand: Text):
        super().__init__('hand_control_node')
        self.app = QApplication(sys.argv)
        self.window = QWidget()
        self.window.resize(800, 600)
        self.layout = QVBoxLayout()
        self.finger_widgets = {}
        self.state_subscriber = self.create_subscription(JointState, f"{hand}/current_joint_states",
                                                         self.listener_callback, 10)
        fingers = ['thumb', 'index', 'middle', 'ring', 'little']
        joints = ['swing', 'rotation']

        for finger in fingers:
            for joint in joints:
                if finger != "thumb" and joint == "swing":
                    continue
                if f"{finger}_{joint}" in self.finger_widgets.keys():
                    continue
                control_widget = FingerControlWidget(self, finger, joint, hand)
                print(f"Adding {finger}_{joint} widget.")
                self.finger_widgets[f"{finger}_{joint}"] = control_widget
                self.layout.addWidget(control_widget)
            # Adds a separator line after thumb and index fingers
            if finger == 'thumb' or finger == 'index' or finger == 'middle':
                line = QFrame()
                line.setFrameShape(QFrame.HLine)
                line.setFrameShadow(QFrame.Plain)
                # Increase the border thickness
                line.setStyleSheet("border: 3px solid gray;")
                line.setStyleSheet("color: gray;")
                self.layout.addWidget(line)
        # Adds text box for initial joint positions
        self.init_pose_textbox = QLineEdit()
        self.init_pose_textbox.setPlaceholderText(
            "Enter initial joint positions as key,value pairs")
        self.layout.addWidget(self.init_pose_textbox)
        # Adds initialization button
        self.init_button = QPushButton("Initialize")
        self.init_button.clicked.connect(self.initialize_joints)
        self.layout.addWidget(self.init_button)

        # Adds aspiration test button
        self.aspiration_button = QPushButton("Aspiration Test")
        self.aspiration_button.clicked.connect(self.aspiration_test)
        self.layout.addWidget(self.aspiration_button)

        self.window.setLayout(self.layout)
        self.window.setWindowTitle("Dexterous Hand Control")
        self.window.show()

    def initialize_joints(self):
        # Index_swing: -0.13
        text = self.init_pose_textbox.text()
        print(text)
        # Text empty, reset all joints to 0
        if not text or len(text) == 0:
            for widget in self.finger_widgets.values():
                widget.publish_value(0.)
            self.ros_spin(dt=0.1)
            return
        for item in text.split(','):
            key, value = item.split(':')
            key = key.strip()
            value = float(value.strip())
            self.finger_widgets[key].set_value(value)
        # Do not clear the text box content

    def aspiration_test(self):
        # Define the series of actions for the aspiration test
        actions = {
            # "thumb_swing": [0, 0.62],
            # "thumb_bend_1": [0, 0.03],
            "thumb_bend_2": [0, 0.6],
            # "index_bend_1": [0, 0.26],
            # "index_bend_2": [0.4, 0.5],
            "middle_bend_1": [0.3, 0.4],
            "middle_bend_2": [0.4, 0.8], }
        # Execute the actions by interpolating the values from start to end
        for i in range(100):
            for key, value in actions.items():
                start, end = value
                self.finger_widgets[key].set_value(
                    start + i * (end - start) / 100)
            self.ros_spin(dt=0.02)

    def listener_callback(self, joint_state):
        print(f'Sub: {joint_state}')
        names = joint_state.name
        for name, value in zip(names, joint_state.position):
            self.finger_widgets[name].set_value(value)
            print(f"current joint: {name}: {value}")

    def run(self):
        timer = self.create_timer(0.01, self.ros_spin)
        sys.exit(self.app.exec())

    def ros_spin(self, dt=0.01):
        rclpy.spin_once(self, timeout_sec=dt)


def main(args=None):
    rclpy.init(args=args)
    node = HandControlNode("rohand_node")
    node.run()
    node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
